import cv2
import numpy as np
import mss
from time import sleep

# 加载 YOLOv3-tiny 模型
net = cv2.dnn.readNet("data/yolov3.weights", "data/yolov3.cfg")

# 使用 GPU 加速
# net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
# net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

# 加载类别名文件（coco.names）
with open("data/coco.names", "r",encoding='gbk') as f:
    classes = [line.strip() for line in f.readlines()]

# 获取 YOLO 输出层的名字
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]

# 定义要检测的屏幕区域坐标：左上角(0, 66)，右下角(420, 980)
x_start, y_start = 0, 66
x_end, y_end = 420, 980
width = x_end - x_start
height = y_end - y_start

# 使用 mss 库来捕获屏幕
with mss.mss() as sct:
    frame_count = 0
    detection_interval = 10 # 每 5 帧进行一次检测

    while True:
        # 捕获指定区域的屏幕
        monitor = {"top": y_start, "left": x_start, "width": width, "height": height}
        screen_img = sct.grab(monitor)

        # 将捕获的屏幕图像转换为 NumPy 数组并从 BGRA 转换为 BGR（OpenCV 使用 BGR）
        frame = np.array(screen_img)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)

        frame_count += 1

        if frame_count % detection_interval == 0:
            # 创建 4D blob
            blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
            net.setInput(blob)

            # 前向传播，获取检测结果
            outs = net.forward(output_layers)

            # 储存检测结果
            class_ids = []
            confidences = []
            boxes = []

            for out in outs:
                for detection in out:
                    scores = detection[5:]
                    class_id = np.argmax(scores)
                    confidence = scores[class_id]

                    # 仅当置信度大于0.5时处理该检测框
                    if confidence > 0.5:
                        center_x = int(detection[0] * width)
                        center_y = int(detection[1] * height)
                        w = int(detection[2] * width)
                        h = int(detection[3] * height)

                        x = int(center_x - w / 2)
                        y = int(center_y - h / 2)

                        boxes.append([x, y, w, h])
                        confidences.append(float(confidence))
                        class_ids.append(class_id)

            # 使用非极大值抑制来去除重叠的检测框
            indexes = cv2.dnn.NMSBoxes(boxes, confidences, score_threshold=0.5, nms_threshold=0.4)

            if len(indexes) > 0:
                for i in indexes.flatten():
                    x, y, w, h = boxes[i]
                    label = str(classes[class_ids[i]])
                    confidence = confidences[i]
                    color = (0, 255, 0)  # 绿色框

                    # 打印识别结果
                    print(f"Detected {label} with confidence {confidence:.2f} at position (x: {x}, y: {y}, w: {w}, h: {h})")

                    # 绘制矩形框
                    cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                    cv2.putText(frame, f"{label} {confidence:.2f}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        # 显示结果，只创建一个窗口
        cv2.imshow("YOLOv3-tiny Screen Region Detection", frame)

        # 按 'q' 键退出，并确保只销毁一个窗口
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        if cv2.waitKey(1) & 0xFF == ord('f'):
            while True:
                sleep(1)
                if cv2.waitKey(1) & 0xFF == ord('g'):
                    break


# 释放窗口资源
cv2.destroyAllWindows()
